import cvxopt
import numpy as np
from ..robot_state import RobotState, LegState
from ..utils.transformation import Transformation

mpc_N = 10
ch = 3
nx = 12
nu = 13

ncfr_single = 4
ncfr = ncfr_single * 2

ncstxya = 1
ncstxy_single = ncstxya * 4
ncstxy = ncstxy_single * 2

ncstza = 2
ncstz_single = ncstza * 4
ncstz = ncstz_single * 2
nc = ncfr + ncstxy + ncstz


class MPC:
    def __init__(self, dt: float) -> None:
        super().__init__()

        self.EN = False

        self.Ac = [np.zeros((nx, nx)) for _ in range(mpc_N)]
        self.A = [np.zeros((nx, nx)) for _ in range(mpc_N)]
        self.Bc = [np.zeros((nx, nu)) for _ in range(mpc_N)]
        self.B = [np.zeros((nx, nu)) for _ in range(mpc_N)]
        self.Cc = np.zeros(nx)
        self.C = np.zeros(nx)

        self.Aqp = np.zeros((nx * mpc_N, nx))
        self.Aqp1 = np.zeros((nx * mpc_N, nx * mpc_N))
        self.Bqp1 = np.zeros((nx * mpc_N, nu * mpc_N))
        self.Bqp = np.zeros((nx * mpc_N, nu * ch))
        self.Cqp1 = np.zeros(nx * mpc_N)
        self.Cqp = np.zeros(nx * mpc_N)

        self.Ufe = np.zeros(nu * ch)
        self.Ufe_pre = np.zeros(nu)
        self.Xd = np.zeros(nx * mpc_N)
        self.X_cur = np.zeros(nx)
        self.X_cal = np.zeros(nx)
        self.X_cal_pre = np.zeros(nx)
        self.dX_cal = np.zeros(nx)

        self.L = np.zeros((nx * mpc_N, nx * mpc_N))
        self.K = np.zeros((nu * ch, nu * ch))
        self.alpha = 0.0
        self.H = np.zeros((nu * ch, nu * ch))
        self.c = np.zeros(nu * ch)

        self.u_low = np.zeros(nu * ch)
        self.u_up = np.zeros(nu * ch)
        self.As = np.zeros((nc * ch, nu * ch))
        self.bs = np.zeros(nc * ch)

        self.m = 59.215542
        self.g = -9.8
        self.miu = 0.5
        self.delta_foot = [0.0 for _ in range(4)]
        self.delta_foot[0] = 0.073
        self.delta_foot[1] = 0.125
        self.delta_foot[2] = 0.025
        self.delta_foot[3] = 0.025

        self.max = [0.0 for _ in range(6)]
        self.min = [0.0 for _ in range(6)]
        self.max[0] = 1000.0
        self.max[1] = 1000.0
        self.max[2] = -3.0 * self.m * self.g
        self.max[3] = 20.0
        self.max[4] = 80.0
        self.max[5] = 100.0
        self.min[0] = -1000.0
        self.min[1] = -1000.0
        self.min[2] = 0.0
        self.min[3] = -20.0
        self.min[4] = -80.0
        self.min[5] = -100.0

        self.pCoM = np.zeros(3)
        self.pf2com = np.zeros(6)
        self.pf2comd = np.zeros(6)
        self.pe = np.zeros(6)
        self.pf2comi = [np.zeros(6) for _ in range(mpc_N)]
        self.Ic = np.zeros((3, 3))
        self.R_curz = [np.zeros((3, 3)) for _ in range(mpc_N)]
        self.R_cur = np.zeros((3, 3))
        self.R_w2f = np.zeros((3, 3))
        self.R_f2w = np.zeros((3, 3))

        self.legStateCur = 0
        self.legStateNext = 0
        self.legState = [0 for _ in range(10)]
        self.dt = dt

        self.qp_H = np.zeros((nu * ch, nu * ch))
        self.qp_As = np.zeros((nc * ch, nu * ch))
        self.qp_c = np.zeros(nu * ch)
        self.qp_lbA = np.zeros(nc * ch)
        self.qp_ubA = np.zeros(nc * ch)
        self.qp_lu = np.zeros(nu * ch)
        self.qp_uu = np.zeros(nu * ch)
        self.nWSR = 100
        self.cpu_time = 0.1
        self.xOpt_iniGuess = np.zeros(nu * ch)

        self.qp_cpuTime = 0.0
        self.qp_Status = 0
        self.qp_nWSR = 0.0

    def set_weight(self, u_weight: float, L_diag: np.ndarray, K_diag: np.ndarray):
        L_diag_N = np.zeros(nx * mpc_N)
        K_diag_N = np.zeros(nu * ch)

        self.L[:, :] = np.zeros((nx * mpc_N, nx * mpc_N))
        self.K[:, :] = np.zeros((nu * ch, nu * ch))

        self.alpha = u_weight
        for i in range(mpc_N):
            L_diag_N[i * nx: (i + 1) * nx] = L_diag

        for i in range(ch):
            K_diag_N[i * nu: (i + 1) * nu] = K_diag

        for i in range(nx * mpc_N):
            self.L[i, i] = L_diag_N[i]

        for i in range(nu * ch):
            self.K[i, i] = K_diag_N[i]

        for i in range(mpc_N):
            self.L[i * nx + 3: i * nx + 6, i * nx + 3: i * nx + 6] = self.R_curz[i] @ self.L[i * nx + 3: i * nx + 6,
                                                                                      i * nx + 3: i * nx + 6] @ \
                                                                     self.R_curz[i].T
            self.L[i * nx + 6: i * nx + 9, i * nx + 6: i * nx + 9] = self.R_curz[i] @ self.L[i * nx + 6: i * nx + 9,
                                                                                      i * nx + 6: i * nx + 9] @ \
                                                                     self.R_curz[i].T
            self.L[i * nx + 9: i * nx + 12, i * nx + 9: i * nx + 12] = self.R_curz[i] @ self.L[i * nx + 9: i * nx + 12,
                                                                                        i * nx + 9: i * nx + 12] @ \
                                                                       self.R_curz[i].T

        for i in range(ch):
            self.K[i * nu: i * nu + 3, i * nu: i * nu + 3] = self.R_curz[i] @ self.K[i * nu: i * nu + 3,
                                                                              i * nu: i * nu + 3] @ self.R_curz[i].T
            self.K[i * nu + 3: i * nu + 6, i * nu + 3: i * nu + 6] = self.R_curz[i] @ self.K[i * nu + 3: i * nu + 6,
                                                                                      i * nu + 3: i * nu + 6] @ \
                                                                     self.R_curz[i].T
            self.K[i * nu + 6: i * nu + 9, i * nu + 6: i * nu + 9] = self.R_curz[i] @ self.K[i * nu + 6: i * nu + 9,
                                                                                      i * nu + 6: i * nu + 9] @ \
                                                                     self.R_curz[i].T
            self.K[i * nu + 9: i * nu + 12, i * nu + 9: i * nu + 12] = self.R_curz[i] @ self.K[i * nu + 9: i * nu + 12,
                                                                                        i * nu + 9: i * nu + 12] @ \
                                                                       self.R_curz[i].T

    def calculate(self):
        if self.EN:
            for i in range(mpc_N):
                self.Ac[i][:3, 6: 9] = self.R_curz[i].transpose()
                self.Ac[i][3: 6, 9: 12] = np.ones(3)
                self.A[i][:, :] = np.eye(nx) + self.dt * self.Ac[i]

            for i in range(mpc_N):
                self.pf2comi[i][:] = self.pf2com
                Ic_W_inv = np.linalg.inv(self.R_curz[i] @ self.Ic @ self.R_curz[i].T)
                self.Bc[i][6: 9, : 3] = Ic_W_inv @ Transformation.CrossProduct_A(self.pf2comi[i][:3])
                self.Bc[i][6: 9, 3: 6] = Ic_W_inv
                self.Bc[i][6: 9, 6: 9] = Ic_W_inv @ Transformation.CrossProduct_A(self.pf2comi[i][3: 6])
                self.Bc[i][6: 9, 9: 12] = Ic_W_inv
                self.Bc[i][9: 12, : 3] = np.eye(3) / self.m
                self.Bc[i][9: 12, 6: 9] = np.eye(3) / self.m
                self.Bc[i][nx - 1, nu - 1] = 1.0 / self.m
                self.B[i] = self.dt * self.Bc[i]

            for i in range(mpc_N):
                self.Aqp[i * nx: (i + 1) * nx, : nx] = np.eye(nx)
            for i in range(mpc_N):
                for j in range(i + 1):
                    self.Aqp[i * nx: (i + 1) * nx, : nx] = self.A[j] @ self.Aqp[i * nx: (i + 1) * nx, : nx]

            for i in range(mpc_N):
                for j in range(i + 1):
                    self.Aqp1[i * nx: (i + 1) * nx, j * nx: (j + 1) * nx] = np.eye(nx)

            for i in range(mpc_N):
                for j in range(i):
                    for k in range(j + 1, i + 1):
                        self.Aqp1[i * nx: (i + 1) * nx, j * nx: (j + 1) * nx] = self.A[k] @ self.Aqp1[
                                                                                            i * nx: (i + 1) * nx,
                                                                                            j * nx: (j + 1) * nx]

            for i in range(mpc_N):
                self.Bqp1[i * nx: (i + 1) * nx, i * nu: (i + 1) * nu] = self.B[i]
            Bqp11 = np.zeros((nu * mpc_N, nu * ch))
            Bqp11[: nu * ch, : nu * ch] = np.eye(nu * ch)
            for i in range(mpc_N - ch):
                Bqp11[nu * ch + i * nu: nu * ch + (i + 1) * nu, nu * (ch - 1): nu * ch] = np.eye(nu)

            B_tmp = np.zeros((nx * mpc_N, nu * ch))
            B_tmp[:, :] = self.Bqp1 @ Bqp11
            self.Bqp[:, :] = self.Aqp1 @ B_tmp

            delta_U = np.zeros(nu * ch)
            for i in range(ch):
                if self.legState[i] == LegState.LST:
                    delta_U[nu * i + 2] = self.m * self.g
                elif self.legState[i] == LegState.RST:
                    delta_U[nu * i + 8] = self.m * self.g
                else:
                    delta_U[nu * i + 2] = 0.5 * self.m * self.g
                    delta_U[nu * i + 8] = 0.5 * self.m * self.g

            self.H[:, :] = 2 * (self.Bqp.T @ self.L @ self.Bqp + self.alpha * self.K) + 1e-10 * np.eye(nu * ch)
            self.c[:] = 2 * self.Bqp.T @ self.L @ (self.Aqp @ self.X_cur - self.Xd) + 0 * self.alpha * self.K @ delta_U

            Asfr111 = np.zeros((ncfr_single, 3))
            Asfr11 = np.zeros((ncfr_single, 3))
            Asfr1 = np.zeros((ncfr, nu))
            Asfr = np.zeros((ncfr * ch, nu * ch))

            Asfr111[:, :] = np.array([
                [-1.0, 0.0, -1.0 / np.sqrt(2.0) * self.miu],
                [1.0, 0.0, -1.0 / np.sqrt(2.0) * self.miu],
                [0.0, -1.0, -1.0 / np.sqrt(2.0) * self.miu],
                [0.0, 1.0, -1.0 / np.sqrt(2.0) * self.miu]
            ])
            Asfr11[:, :] = Asfr111 @ self.R_w2f
            Asfr1[0: ncfr_single, : 3] = Asfr11
            Asfr1[ncfr_single: 2 * ncfr_single, 6: 9] = Asfr11

            for i in range(ch):
                Asfr[ncfr * i: ncfr * (i + 1), nu * i: nu * (i + 1)] = Asfr1

            sign_xy = [1.0, -1.0, -1.0, 1.0]
            gxyz = [np.zeros(3) for _ in range(4)]
            gxyz[0][:] = np.array([0.0, 1.0, 0.0])
            gxyz[1][:] = np.array([0.0, 1.0, 0.0])
            gxyz[2][:] = np.array([1.0, 0.0, 0.0])
            gxyz[3][:] = np.array([1.0, 0.0, 0.0])
            r = [np.zeros(3) for _ in range(4)]
            p = [np.zeros(3) for _ in range(4)]
            Astxy_r = [np.zeros((ncstxya, 6)) for _ in range(4)]
            Astxy11 = np.zeros((ncstxy_single, 6))
            Astxy1 = np.zeros((ncstxy, nu))
            Astxy = np.zeros((ncstxy * ch, nu * ch))

            r[0][:] = np.array([0.0, 1.0, 0.0])
            r[1][:] = np.array([0.0, 1.0, 0.0])
            r[2][:] = np.array([1.0, 0.0, 0.0])
            r[3][:] = np.array([1.0, 0.0, 0.0])

            p[0][:] = np.array([self.delta_foot[0], 0.0, 0.0])
            p[1][:] = np.array([-self.delta_foot[1], 0.0, 0.0])
            p[2][:] = np.array([0.0, self.delta_foot[2], 0.0])
            p[3][:] = np.array([0.0, -self.delta_foot[3], 0.0])

            for i in range(4):
                Astxy_r[i][0, : 3] = sign_xy[i] * gxyz[i] @ r[i] * np.cross(self.R_f2w @ r[i], self.R_f2w @ p[i])
                Astxy_r[i][0, 3: 6] = sign_xy[i] * gxyz[i] @ self.R_w2f
                Astxy11[i * ncstxya: (i + 1) * ncstxya, : 6] = Astxy_r[i]
            Astxy1[: ncstxy_single, : 6] = Astxy11
            Astxy1[ncstxy_single: 2 * ncstxy_single, 6: 12] = Astxy11
            for i in range(ch):
                Astxy[ncstxy * i: ncstxy * (i + 1), nu * i: nu * (i + 1)] = Astxy1

            Astz_r = [np.zeros((ncstza, 6)) for _ in range(4)]
            Astz11 = np.zeros((ncstz_single, 6))
            Astz1 = np.zeros((ncstz, nu))
            Astz = np.zeros((ncstz * ch, nu * ch))

            for i in range(4):
                Astz_r[i][0, :3] = - np.sqrt(
                    p[i][0] * p[i][0] + p[i][1] * p[i][1] + p[i][2] * p[i][2]) * self.miu * self.R_w2f[2, :]
                Astz_r[i][0, 3: 6] = self.R_w2f[2, :]
                Astz_r[i][1, :3] = Astz_r[i][0, : 3]
                Astz_r[i][1, 3: 6] = -Astz_r[i][0, 3: 6]
                Astz11[i * ncstza: (i + 1) * ncstza, : 6] = Astz_r[i]

            Astz1[0: ncstz_single, : 6] = Astz11
            Astz1[ncstz_single: 2 * ncstz_single, 6: 12] = Astz11
            for i in range(ch):
                Astz[ncstz * i: ncstz * (i + 1), nu * i: nu * (i + 1)] = Astz1

            self.As[: ncfr * ch, : nu * ch] = Asfr
            self.As[ncfr * ch: ncfr * ch + ncstxy * ch, : nu * ch] = Astxy
            self.As[ncfr * ch + ncstxy * ch: ncfr * ch + ncstxy * ch + ncstz * ch, : nu * ch] = Astz
            self.bs[:] *= 0

            Guess_value = np.zeros(nu * ch)
            for i in range(ch):
                if self.legState[i] == LegState.DST:
                    Guess_value[i * nu + 2] = -0.5 * self.m * self.g
                    Guess_value[i * nu + 8] = -0.5 * self.m * self.g
                    Guess_value[i * nu + 12] = self.m * self.g
                    for j in range(6):
                        self.u_low[i * nu + j] = self.min[j]
                        self.u_low[i * nu + j + 6] = self.min[j]
                        self.u_up[i * nu + j] = self.max[j]
                        self.u_up[i * nu + j + 6] = self.max[j]
                    self.u_low[i * nu + 12] = self.m * self.g
                    self.u_up[i * nu + 12] = self.m * self.g
                elif self.legState[i] == LegState.LST:
                    Guess_value[i * nu + 2] = -self.m * self.g
                    Guess_value[i * nu + 8] = 0.0
                    Guess_value[i * nu + 12] = self.m * self.g
                    for j in range(6):
                        self.u_low[i * nu + j] = self.min[j]
                        self.u_low[i * nu + j + nu // 2] = 0.0
                        self.u_up[i * nu + j] = self.max[j]
                        self.u_up[i * nu + j + nu // 2] = 0.0
                    self.u_low[i * nu + 12] = self.m * self.g
                    self.u_up[i * nu + 12] = self.m * self.g
                elif self.legState[i] == LegState.RST:
                    Guess_value[i * nu + 2] = 0.0
                    Guess_value[i * nu + 8] = -self.m * self.g
                    Guess_value[i * nu + 12] = self.m * self.g
                    for j in range(6):
                        self.u_low[i * nu + j] = 0.0
                        self.u_low[i * nu + j + nu // 2] = self.min[j]
                        self.u_up[i * nu + j] = 0.0
                        self.u_up[i * nu + j + nu // 2] = self.max[j]
                    self.u_low[i * nu + 12] = self.m * self.g
                    self.u_up[i * nu + 12] = self.m * self.g

            self.nWSR = 1000000
            self.cpu_time = self.dt

            lbA = np.zeros(nc * ch)
            ubA = np.zeros(nc * ch)
            one_ch_1 = np.ones(nc * ch)
            lbA = -1e7 * one_ch_1
            ubA = 1e7 * one_ch_1

            for i in range(ch):
                if self.legState[i] == LegState.DST:
                    ubA[ncfr * i: ncfr * (i + 1)] *= 0
                    ubA[ncfr * ch + ncstxy * i: ncfr * ch + ncstxy * (i + 1)] *= 0
                    ubA[ncfr * ch + ncstxy * ch + ncstz * i: ncfr * ch + ncstxy * ch + ncstz * (i + 1)] = 0
                elif self.legState[i] == LegState.LST:
                    ubA[ncfr * i: ncfr * i + ncfr_single] *= 0
                    ubA[ncfr * ch + ncstxy * i: ncfr * ch + ncstxy * i + ncstxy_single] *= 0
                    ubA[ncfr * ch + ncstxy * ch + ncstz * i: ncfr * ch + ncstxy * ch + ncstz * i + ncstz_single] = 0
                elif self.legState[i] == LegState.RST:
                    ubA[ncfr * i + ncfr_single: ncfr * i + ncfr_single * 2] *= 0
                    ubA[ncfr * ch + ncstxy * i + ncstxy_single: ncfr * ch + ncstxy * i + ncstxy_single * 2] *= 0
                    ubA[
                    ncfr * ch + ncstxy * ch + ncstz * i + ncstz_single: ncfr * ch + ncstxy * ch + ncstz * i + ncstz_single * 2] = 0

            self.qp_H = cvxopt.matrix(self.H)
            self.qp_c = cvxopt.matrix(self.c)
            Au = np.eye(nu * ch)
            self.qp_As = cvxopt.matrix(np.vstack((self.As, -self.As, Au, -Au)))
            self.qp_lbA = cvxopt.matrix(lbA)
            self.qp_ubA = cvxopt.matrix(ubA)
            self.qp_lu = cvxopt.matrix(self.u_low)
            self.qp_uu = cvxopt.matrix(self.u_up)
            b = cvxopt.matrix(np.hstack((ubA, -lbA, self.u_up, -self.u_low)))

            res = cvxopt.solvers.qp(self.qp_H, self.qp_c, self.qp_As, b)
            self.qp_nWSR = self.nWSR
            self.qp_cpuTime = self.cpu_time

            x_opt = res['x']
            for i in range(nu * ch):
                self.Ufe[i] = x_opt[i]

            self.dX_cal = self.Ac[0] @ self.X_cur + self.Bc[0] @ self.Ufe[: nu]
            delta_X = np.zeros(nx)
            for i in range(3):
                delta_X[i] = 0.5 * self.dX_cal[i + 6] * self.dt * self.dt
                delta_X[i + 3] = 0.5 * self.dX_cal[i + 9] * self.dt * self.dt
                delta_X[i + 6] = self.dX_cal[i + 6] * self.dt
                delta_X[i + 9] = self.dX_cal[i + 9] * self.dt

            self.X_cal = (self.Aqp @ self.X_cur + self.Bqp @ self.Ufe)[: nx] + delta_X
            self.Ufe_pre = self.Ufe[: nu]

    def read_state(self, robot_state: RobotState):
        self.X_cur[:3] = robot_state.base_rpy
        self.X_cur[3: 6] = robot_state.q[0: 3]
        self.X_cur[6: 9] = robot_state.dq[3: 6]
        self.X_cur[9: 12] = robot_state.dq[: 3]

        if self.EN:
            for i in range(mpc_N - 1):
                self.Xd[nx * i: nx * (i + 1)] = self.Xd[nx * (i + 1): nx * (i + 2)]
            for j in range(3):
                self.Xd[nx * (mpc_N - 1) + j] = robot_state.js_eul_des[j]
                self.Xd[nx * (mpc_N - 1) + 3 + j] = robot_state.js_pos_des[j]
                self.Xd[nx * (mpc_N - 1) + 6 + j] = robot_state.js_omega_des[j]
                self.Xd[nx * (mpc_N - 1) + 9 + j] = robot_state.js_vel_des[j]

        else:
            for i in range(mpc_N):
                for j in range(3):
                    self.Xd[nx * i + j] = self.X_cur[j]
                    self.Xd[nx * i + 3 + j] = self.X_cur[3 + j]
                    self.Xd[nx * i + 6 + j] = self.X_cur[6 + j]
                    self.Xd[nx * i + 9 + j] = self.X_cur[9 + j]

            for i in range(mpc_N):
                for j in range(3):
                    self.Xd[nx * i + j] = robot_state.js_eul_des[j]
                    self.Xd[nx * i + 3 + j] = robot_state.js_pos_des[j]
                    self.Xd[nx * i + 6 + j] = robot_state.js_omega_des[j]
                    self.Xd[nx * i + 9 + j] = robot_state.js_vel_des[j]

        self.R_cur[:, :] = Transformation.eul2Rot(*self.X_cur[: 3])
        for i in range(mpc_N):
            self.R_curz[i] = Transformation.Rz3(self.X_cur[2])
        self.pCoM[:] = self.X_cur[3: 6]
        self.pe[:3] = robot_state.fe_l_pos_W
        self.pe[3: 6] = robot_state.fe_r_pos_W

        self.pf2com[: 3] = self.pe[:3] - self.pCoM
        self.pf2com[3: 6] = self.pe[3: 6] - self.pCoM
        self.pf2comd[: 3] = self.pe[: 3] - self.Xd[3: 6]
        self.pf2comd[3: 6] = self.pe[3: 6] - self.Xd[3: 6]

        self.Ic[:, :] = robot_state.inertia

        self.legStateCur = robot_state.legState
        self.legStateNext = robot_state.legStateNext
        for i in range(mpc_N):
            aa = i * self.dt / 0.2
            phip = robot_state.phi + aa
            if phip > 1:
                self.legState[i] = self.legStateNext
            else:
                self.legState[i] = self.legStateCur

        R_slop = np.zeros((3, 3))
        R_slop[:, :] = Transformation.eul2Rot(*robot_state.slop)
        if self.legStateCur == LegState.RST:
            self.R_f2w[:, :] = robot_state.fe_r_rot_W
        elif self.legStateCur == LegState.LST:
            self.R_f2w[:, :] = robot_state.fe_l_rot_W
        else:
            self.R_f2w[:, :] = R_slop
        self.R_w2f[:, :] = self.R_f2w.T

    def write_state(self, robot_state: RobotState):
        robot_state.Xd[:] = self.Xd
        robot_state.X_cur[:] = self.X_cur
        robot_state.fe_react_tau_cmd[:] = self.Ufe[:12]
        robot_state.X_cal[:] = self.X_cal
        robot_state.dX_cal[:] = self.dX_cal

        robot_state.qpStatus_MPC = self.nWSR
        robot_state.qp_cpuTime_MPC = self.cpu_time
        robot_state.qpStatus_MPC = self.qp_Status

        robot_state.Fr_ff[:] = self.Ufe[:12]

        k = 5
        robot_state.des_ddq[:2] = self.dX_cal[9: 11]
        robot_state.des_ddq[5] = k * (self.Xd[6 + 2] - robot_state.dq[5])

        robot_state.des_dq[:3] = self.Xd[9: 12]
        robot_state.des_dq[3: 5] = np.zeros(2)
        robot_state.des_dq[5] = self.Xd[6 + 2]

        robot_state.des_delta_q[: 2] = robot_state.des_dq[: 2] * self.dt
        robot_state.des_delta_q[5] = robot_state.des_dq[5] * self.dt

        robot_state.base_rpy_des = np.array([0.005, 0.0, self.Xd[2]])
        robot_state.base_pos_des = self.Xd[3: 6]

    def enable(self):
        self.EN = True

    def disable(self):
        self.EN = False

    def get_EN(self):
        return self.EN
